package com.education.platform.websocket;

import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.education.platform.service.MessageService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 聊天WebSocket处理器
 *
 * @author education-platform
 * @since 2024-10-07
 */
@Component
@RequiredArgsConstructor
@Slf4j
public class ChatWebSocketHandler implements WebSocketHandler {

    private final MessageService messageService;
    
    // 存储用户ID与WebSocket会话的映射
    private static final Map<Long, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();
    
    // 存储会话ID与用户ID的映射
    private static final Map<String, Long> SESSION_USERS = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long userId = (Long) session.getAttributes().get("userId");
        if (userId != null) {
            USER_SESSIONS.put(userId, session);
            SESSION_USERS.put(session.getId(), userId);
            log.info("用户 {} 建立WebSocket连接, 会话ID: {}", userId, session.getId());
            
            // 发送连接成功消息
            sendMessage(session, createMessage("CONNECT", "连接成功", null));
            
            // 通知用户上线
            messageService.handleUserOnline(userId);
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        Long userId = SESSION_USERS.get(session.getId());
        if (userId == null) {
            log.warn("未找到会话对应的用户ID: {}", session.getId());
            return;
        }
        
        try {
            String payload = message.getPayload().toString();
            JSONObject jsonMessage = JSONUtil.parseObj(payload);
            
            String type = jsonMessage.getStr("type");
            log.info("收到用户 {} 的消息, 类型: {}, 内容: {}", userId, type, payload);
            
            switch (type) {
                case "CHAT":
                    handleChatMessage(userId, jsonMessage);
                    break;
                case "READ":
                    handleReadMessage(userId, jsonMessage);
                    break;
                case "TYPING":
                    handleTypingMessage(userId, jsonMessage);
                    break;
                case "HEARTBEAT":
                    handleHeartbeat(session);
                    break;
                default:
                    log.warn("未知消息类型: {}", type);
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息异常", e);
            sendMessage(session, createMessage("ERROR", "消息处理失败", null));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        Long userId = SESSION_USERS.get(session.getId());
        log.error("WebSocket传输异常, 用户: {}, 会话: {}", userId, session.getId(), exception);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        Long userId = SESSION_USERS.remove(session.getId());
        if (userId != null) {
            USER_SESSIONS.remove(userId);
            log.info("用户 {} 断开WebSocket连接, 会话ID: {}, 状态: {}", userId, session.getId(), closeStatus);
            
            // 通知用户下线
            messageService.handleUserOffline(userId);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 处理聊天消息
     */
    private void handleChatMessage(Long fromUserId, JSONObject message) {
        try {
            Long toUserId = message.getLong("toUserId");
            String content = message.getStr("content");
            String contentType = message.getStr("contentType", "TEXT");
            
            // 保存消息到数据库
            Long messageId = messageService.sendMessage(fromUserId, toUserId, content, contentType, null);
            
            // 给发送者发送确认消息
            WebSocketSession fromSession = USER_SESSIONS.get(fromUserId);
            if (fromSession != null && fromSession.isOpen()) {
                JSONObject ackMessage = createMessage("MESSAGE_SENT", "消息发送成功", JSONUtil.createObj()
                        .set("messageId", messageId)
                        .set("toUserId", toUserId)
                        .set("content", content)
                        .set("contentType", contentType)
                        .set("timestamp", System.currentTimeMillis()));
                
                sendMessage(fromSession, ackMessage);
            }
            
            // 实时推送给接收者
            WebSocketSession toSession = USER_SESSIONS.get(toUserId);
            if (toSession != null && toSession.isOpen()) {
                JSONObject pushMessage = createMessage("NEW_MESSAGE", "新消息", JSONUtil.createObj()
                        .set("messageId", messageId)
                        .set("fromUserId", fromUserId)
                        .set("content", content)
                        .set("contentType", contentType)
                        .set("timestamp", System.currentTimeMillis()));
                
                sendMessage(toSession, pushMessage);
            }
            
        } catch (Exception e) {
            log.error("处理聊天消息异常", e);
            
            // 发送错误消息给发送者
            WebSocketSession fromSession = USER_SESSIONS.get(fromUserId);
            if (fromSession != null && fromSession.isOpen()) {
                JSONObject errorMessage = createMessage("MESSAGE_ERROR", "消息发送失败", JSONUtil.createObj()
                        .set("error", e.getMessage()));
                sendMessage(fromSession, errorMessage);
            }
        }
    }

    /**
     * 处理消息已读
     */
    private void handleReadMessage(Long userId, JSONObject message) {
        try {
            Long messageId = message.getLong("messageId");
            messageService.markMessageAsRead(messageId, userId);
            
        } catch (Exception e) {
            log.error("处理消息已读异常", e);
        }
    }

    /**
     * 处理正在输入
     */
    private void handleTypingMessage(Long fromUserId, JSONObject message) {
        try {
            Long toUserId = message.getLong("toUserId");
            Boolean isTyping = message.getBool("isTyping", false);
            
            // 推送给对方
            WebSocketSession toSession = USER_SESSIONS.get(toUserId);
            if (toSession != null && toSession.isOpen()) {
                JSONObject typingMessage = createMessage("TYPING", "正在输入", JSONUtil.createObj()
                        .set("fromUserId", fromUserId)
                        .set("isTyping", isTyping));
                
                sendMessage(toSession, typingMessage);
            }
            
        } catch (Exception e) {
            log.error("处理正在输入消息异常", e);
        }
    }

    /**
     * 处理心跳
     */
    private void handleHeartbeat(WebSocketSession session) {
        sendMessage(session, createMessage("HEARTBEAT", "pong", null));
    }

    /**
     * 发送消息给指定用户
     */
    public void sendMessageToUser(Long userId, String type, String message, Object data) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        if (session != null && session.isOpen()) {
            sendMessage(session, createMessage(type, message, data));
        }
    }

    /**
     * 广播消息给所有在线用户
     */
    public void broadcastMessage(String type, String message, Object data) {
        JSONObject broadcastMessage = createMessage(type, message, data);
        USER_SESSIONS.values().forEach(session -> {
            if (session.isOpen()) {
                sendMessage(session, broadcastMessage);
            }
        });
    }

    /**
     * 获取在线用户列表
     */
    public java.util.Set<Long> getOnlineUsers() {
        return USER_SESSIONS.keySet();
    }

    /**
     * 检查用户是否在线
     */
    public boolean isUserOnline(Long userId) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        return session != null && session.isOpen();
    }

    /**
     * 发送消息
     */
    private void sendMessage(WebSocketSession session, JSONObject message) {
        try {
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(message.toString()));
            }
        } catch (IOException e) {
            log.error("发送WebSocket消息失败", e);
        }
    }

    /**
     * 创建消息对象
     */
    private JSONObject createMessage(String type, String message, Object data) {
        JSONObject result = JSONUtil.createObj()
                .set("type", type)
                .set("message", message)
                .set("timestamp", System.currentTimeMillis());
        
        if (data != null) {
            result.set("data", data);
        }
        
        return result;
    }
}
